'''
ipV4 Geographical Lookup
Written by: Julian Tosh

This module required ipV4address.py (see my other projects).

This module reads in a list of CIDR blocks with associations with countries
and other known locations, such as buildings, sites, subnets, hosts, etc.

A good source for IPs associated with countries, check out: 
	http://www.countryipblocks.net/

'''

import ipV4address
import bisect
import os
import pickle
from math import sqrt,log

class ipV4GeoLookup:
    def __init__(self, geoFile):
        if len(geoFile) == 0 :
            raise RuntimeError, "You must supply a geo table file."

        self.__debugPrint__("Loading " + geoFile + "...")
        self.__loadLocationList__(geoFile)


    def __debugPrint__(self, string):
        print string
        pass
    

    # GeoList format:
    # [0] : IP Address / Subnet
    # [1] : Mask
    # [2] : Broadcast Address
    # [3] : 2 letter ISO country code
    # [4] : ISO country
    # [5] : free form additional description
    #geoList = [  []  ]
    geoList = []


    def __loadLocationList__(self, txtFileLoc):
        objFileLoc = txtFileLoc + ".obj"     
        self.__debugPrint__("Choosing cache or rebuild...")
        
        #
        # Get timestamp of text file
        #
        if os.path.isfile(txtFileLoc):
            txtFileTimeStamp = os.stat(txtFileLoc).st_mtime
            self.__debugPrint__(str(txtFileTimeStamp) + " is the timestamp on " + txtFileLoc)
        else:
            raise RuntimeError, "Cannot read file " + txtFileLoc + ". Goodbye."
        
        #
        # Get timestamp of object file
        #
        if os.path.isfile(objFileLoc):
            objFileTimeStamp = os.stat(objFileLoc).st_mtime
            self.__debugPrint__(str(objFileTimeStamp) + " is the timestamp on " + objFileLoc)
        else:
            self.__debugPrint__(objFileLoc + " file does not exist.")
            objFileTimeStamp = 0            
        
        if True:   # Debug for caching - always rebuild from text file if true
            self.__debugPrint__("Bypassing timestamp check. Reading text file...")
            self.__rebuildGeoArrayFromFile__(txtFileLoc)        	
        elif txtFileTimeStamp >= objFileTimeStamp:
            self.__debugPrint__("Text file timestamp is newer. Reading text file...")
            self.__rebuildGeoArrayFromFile__(txtFileLoc)
        else:
            self.__debugPrint__("Object file timestamp is newer. Reading obj file...")
            self.__loadGeoArrayFromPickle__(objFileLoc)
        

    def __rebuildGeoArrayFromFile__(self, txtFileLoc):
        self.__debugPrint__("Rebuilding geoArray from text file: " + txtFileLoc)
        
        lineCount = 0
        
        #
        # Open text file
        #        
        try:
            file = open(txtFileLoc, 'r')
        except:
            raise RuntimeError, "Trouble reading specified file: " + txtFileLoc
        
        #
        # process text file
        #
        for line in file:
            lineCount += 1
        
            if line[:11] == "# Country: ":
                title,locationLong = line.split(":")
                locationLong = locationLong.strip()
                continue
            elif line[:12]  ==  "# ISO Code: ":
                title,locationShort = line.split(":")
                locationShort = locationShort.strip()
                continue
            elif line[:1] == "#":            
                continue
            elif len(line.strip()) == 0:
            	continue
            else:
                #
                # Line is not a comment or country line, so it must contain an IP address.
                #            
                
                # 'Extra' field begins after the first space. IP lines are formatted as: [IP][Space][Extra description]
                delim = line.find(" ") #i.e. extra begins after first space. delim is -1 if no space
                #self.__debugPrint__("::::" + str(delim) + "::" + line[:delim] + "::::" + line[delim:])
                
                if delim <= 0:
                    ip = ipV4address.ipaddress(line)
                    extra = ""
                else:
                    ip = ipV4address.ipaddress(line[:delim])
                    extra = line[delim:].strip()
                
                #self.__debugPrint__("Inserting: " + str(ip.ipAsInt) + " - " + str(ip.ipMaskAsInt) + " - " + str(ip.ipBroadcastAsInt) + " - " + locationShort + " - " + locationLong + " - " + extra)
                
                if self.geoList == None:
					#self.geoList = [ [ip.ipAsInt, ip.ipMaskAsInt, ip.ipBroadcastAsInt, locationShort, locationLong, extra] ]
					self.geoList.append( {'ip' : ip.ipAsInt, 'broadcast': ip.ipBroadcastAsInt, 'short' : locationShort, 'long' : locationLong, 'extra' : extra} )
                else:
                	#bisect.insort(self.geoList, [ ip.ipAsInt, ip.ipMaskAsInt, ip.ipBroadcastAsInt, locationShort, locationLong, extra ])
                	bisect.insort(self.geoList,{'ip' : ip.ipAsInt, 'mask': ip.ipMaskAsInt, 'broadcast': ip.ipBroadcastAsInt, 'short' : locationShort, 'long' : locationLong, 'extra' : extra} )

        self.__debugPrint__("Saving array object...")
        self.__saveGeoArrayToPickle__(txtFileLoc)
        
        self.__debugPrint__(str(lineCount) + " lines read. geoList array now has " + str(len(self.geoList)) + " elements. Goodbye.")


    def __loadGeoArrayFromPickle__(self, objFileLoc):
        self.__debugPrint__("Loading geoArray from Pickle...")
        f = open(objFileLoc, 'r')
        self.geoList = pickle.load(f)
        

    def __saveGeoArrayToPickle__(self, txtFileLoc):
        self.__debugPrint__("Saving geoArray to Pickle...")
        objFileLoc = txtFileLoc + ".obj"
        
        f = open(objFileLoc, 'w')
        pickle.dump(self.geoList, f)


    def __printGeoArray2__(self):
 		print "GeoList contains " + str(len(self.geoList)) + " items."
		print "IP".center(8) + " " + "Bcast".center(8) + " " + "Mask".center(8) 
		for d in self.geoList:
			print (  
					hex(d['ip'])[2:-1].upper().rjust(8) + " " + 
					hex(d['broadcast'])[2:-1].upper().rjust(8) + " " +
					hex(d['mask'])[2:-1].upper().rjust(8) + " : " +
					d['short'] + " : " + d['long'] + " : " + d['extra']
              	  )


    def __findMatches__(self, ip, maskBits = 32):
        #
        # Convert doted decimal to integer
        #
        searchIP = ipV4address.ipaddress(ip, maskBits)
        ipInt = searchIP.ipAsInt

        matches = []
            
        #
        # Perform binary search for matching subnet
        #
        indexMax = len(self.geoList)
        #maxLoop = indexMax
        maxLoop = int(log(indexMax, 2)*2)
        self.__debugPrint__("Maximum iterations should be: " + str(maxLoop))
        indexMin = 0
        indexCur = (int(indexMax - indexMin) / 2) + indexMin 
        match = False
        loop = 0
        
        while (match == False):
            loop += 1
            
            #
            # Divide list in have until you find the right one.
            #
            subnet = self.geoList[indexCur]['ip']
            mask = self.geoList[indexCur]['mask']
            broadcast = self.geoList[indexCur]['broadcast']
            match = ((ipInt & mask) == subnet)

            self.__debugPrint__( ( str(loop) + 
                       " Matching IP: " + ip + "/" + hex(ipInt)[2:-1].upper() + 
                       ", Subnet:" + hex(subnet)[2:-1].upper() + 
                       ", Mask:" + hex(mask)[2:-1].upper() + 
                       ", " + str(match).ljust(6) + ", " + 
                       self.geoList[indexCur]['short']
                                  ) )
            
            if match:
                self.__debugPrint__("Match found. Break!")
                #
                # TEST: Checking next record for match
                #
                testOffset = -1
                subnet2 = self.geoList[indexCur + testOffset]['ip']
                mask2= self.geoList[indexCur + testOffset]['mask']
                match2 = (((ipInt & mask2) & mask) == subnet)
                self.__debugPrint__( ( "Before: " + 
                           hex(self.geoList[indexCur + testOffset]['ip'])[2:-1].upper() + " " + 
                           hex(self.geoList[indexCur + testOffset]['mask'])[2:-1].upper() + " " + 
                           hex(self.geoList[indexCur + testOffset]['broadcast'])[2:-1].upper() + " " + 
                           str(match2).ljust(6) + 
                           self.geoList[indexCur + testOffset]['short'] + " " + 
                           self.geoList[indexCur + testOffset]['long'] + " " +
                           self.geoList[indexCur + testOffset]['extra'] + " "
                                      ) ) 

                testOffset = 0
                subnet2 = self.geoList[indexCur + testOffset]['ip']
                mask2= self.geoList[indexCur + testOffset]['mask']
                match2 = (((ipInt & mask2) & mask) == subnet)
                self.__debugPrint__( ( "Match:  " + 
                           hex(self.geoList[indexCur + testOffset]['ip'])[2:-1].upper() + " " + 
                           hex(self.geoList[indexCur + testOffset]['mask'])[2:-1].upper() + " " +
                           hex(self.geoList[indexCur + testOffset]['broadcast'])[2:-1].upper() + " " + 
                           str(match2).ljust(6) + 
                           self.geoList[indexCur + testOffset]['short'] + " " + 
                           self.geoList[indexCur + testOffset]['long'] 
                                     ) ) 

                testOffset = 1
                subnet2 = self.geoList[indexCur + testOffset]['ip']
                mask2= self.geoList[indexCur + testOffset]['mask']
                match2 = (((ipInt & mask2) & mask) == subnet)
                self.__debugPrint__( ( "After:  " + 
                           hex(self.geoList[indexCur + testOffset]['ip'])[2:-1].upper() + " " + 
                           hex(self.geoList[indexCur + testOffset]['mask'])[2:-1].upper() + " " + 
                           hex(self.geoList[indexCur + testOffset]['broadcast'])[2:-1].upper() + " " + 
                           str(match2).ljust(6) + 
                           self.geoList[indexCur + testOffset]['short'] + " " + 
                           self.geoList[indexCur + testOffset]['long'] 
                                     ) ) 
                #
                # TEST: END
                #
                
                break
            elif ipInt > subnet:
                self.__debugPrint__("  Going lower")
                indexMin = indexCur
                indexCur = (int(indexMax - indexMin) / 2) + indexMin 
            else:
                self.__debugPrint__("  Going higher")
                indexMax = indexCur
                indexCur = (int(indexMax - indexMin) / 2) + indexMin
    
            #
            # Race Condition
            #
            if loop > maxLoop:
                self.__debugPrint__("ERROR: Race condition met.")
                return "Unknown"
        
        print "Returning dictionary: " + str(self.geoList[indexCur])            
        return self.geoList[indexCur]




###
### Public Functions
###

    def findHostMatches(self,ip):
        return self.__findMatches__(ip)
    
    
    def findSubnetMatches(self, ip, maskBits):
        pass
    

    def findSupernetMatches(self, ip, maskBits):
    	pass
    
       
    def ipIsIn(self, ip, locationList):
        ipLocation = self.__findMatches__(ip)
        for locationShort in locationList:
            if locationShort == ipLocation:
                return True
        return False








###
### Module Tests
###
def test0():
	'''
	Building array, testing caching features.
	'''
	myGeoList = ipV4GeoLookup("./geoData/subnetTest.txt")
	myGeoList.__printGeoArray2__()

def test1():
	'''
	Finding a host's subnets.
	'''
	myGeoList = ipV4GeoLookup("./geoData/subnetTest.txt")
	ip = "10.10.10.10"
	match = myGeoList.findHostMatches(ip)
	print "Found this IP in " + match['short']









def test2():
    myGeoList = ipV4GeoLookup("./geoCodes/geoList.txt")
    ip = "123.123.123.123"
    sensitiveCountries = ["CN", "RU", "IL"]
    print ip + " is in a sensitive country: " + str(myGeoList.ipIsIn(ip, sensitiveCountries))

def test3():
    myGeoList = ipV4GeoLookup("./geoCodes/all.txt")

    #printArray()
    
    ip = "198.124.5.24"
    location = myGeoList.__findMatches__(ip)
    print "This IP resides in: " + location[0] + " aka " + location[1] 
    
def test4():
    myGeoList = ipV4GeoLookup("./geoCodes/known.txt")

    myGeoList.__printGeoArray__()
    
    ip = "206.206.185.1"
    location = myGeoList.__findMatches__(ip)
    print "This IP resides in: " + location[0] + " aka " + location[1] 
    
def test5():
    myGeoList = ipV4GeoLookup("./geoCodes/known.txt")

    myGeoList.__printGeoArray__()
    
    ip = "159.64.199.1"
    location = myGeoList.__findMatches__(ip)
    print "This IP resides in: " + location[0] + " aka " + location[1] 
    

if __name__ == '__main__':
    test0()
    test1()
    #test2()
    #test3()
    #test4()
    #test5()
    #test6()
    
    